# 结合数据库实现注册和登录功能
import pymysql
import hashlib
import random

# 当前用户
user_info = {'id': 0, 'name': '未登录'}
# 数据库连接
db_conn = pymysql.connect(
    host='124.221.242.12', port=8006,
    user='root', password='123456',
    db='test1', charset='utf8',
    autocommit=True,
)


def encryption(password: str, n: int, salt: str = None) -> tuple[str, str]:
    """密码加密"""
    if salt is None:
        salt = ''
        for i in range(n):
            num = str(random.randint(0, 9))
            upper = chr(random.randint(65, 90))
            lower = chr(random.randint(97, 122))
            random_code = random.choice([num, lower, upper])
            salt = ''.join([salt, random_code])
    m = hashlib.md5(password.encode('utf8'))
    m.update(salt.encode('utf8'))
    return m.hexdigest(), salt


def register():
    """注册"""
    username = input('请输入注册的用户名:').strip()
    password = input('请输入注册的用户密码:').strip()
    re_password = input('请再次输入用户密码:').strip()
    # 密码不一致
    if password != re_password:
        print('两次密码输入不一致！')
        return
    # 密码加盐
    encryption_password, salt = encryption(password, 6)
    # 执行sql语句添加数据
    register_sql = f"insert into user_info(name,password,salt) values('{username}','{encryption_password}','{salt}');"
    cursor = db_conn.cursor()
    affect_rows = cursor.execute(register_sql)
    cursor.close()
    if affect_rows:  # 记录新增成功
        print(f'{username}注册成功！')
        return
        # 记录新增失败
    print(f'{username}注册失败！')


def login():
    """登录"""
    username = input('请输入用户名:').strip()
    password = input('请输入用户密码:').strip()

    login_sql = f"select * from user_info where name = '{username}';"
    cursor = db_conn.cursor()
    affect_rows = cursor.execute(login_sql)
    if not affect_rows:  # 没用记录
        print('当前用户不存在！请先注册！')
        cursor.close()
        return
    # 查询到记录
    for user in cursor.fetchall():
        if user[2] == encryption(password, 6, user[3])[0]:
            print(f'{username}登录成功！')
            user_info['id'] = user[0]
            user_info['name'] = username
            return
    else:
        print('用户名或密码错误！')

    # 选项功能字典


func_dict = {
    '1': register,
    '2': login
}


def run():
    """主循环"""
    while True:
        print(f'当前用户:{user_info["name"]}')
        print('1.注册')
        print('2.登录')
        print('Q退出T登出')
        order = input('请输入选项:').strip()
        # 退出登录
        if order == 'T':
            if user_info['name'] == '未登录':
                print('请先登录！')
                input('按回车继续...')
                continue
            print(f'{user_info["name"]}退出登录！')
            user_info['id'] = 0
            user_info['name'] = '未登录'
            input('按回车继续...')
            continue
        # 退出系统
        if order == 'Q':
            db_conn.close()  # 关闭数据库
            exit(0)
        # 输入错误
        if not func_dict.get(order):
            print('输出错误！请重新输入！')
            input('按回车继续...')
            continue
        # 执行功能
        func_dict[order]()
        input('按回车继续...')


if __name__ == '__main__':
    run()
